#include "dicom/storage/local_cache.h"

#include <assert.h>
#include "dicom/storage/encrypt.h"

namespace dicom {

enum FileCacheState {
  kFileMoving = 1,
  kFileNormal = 2,
  kFileFailed = 4,
};

LRUFileCache::LRUFileCache(const std::string& workdir, int64_t capacity)
    : workdir_(workdir), used_(0), capacity_(capacity) {}

LRUFileCache::~LRUFileCache() {}

bool LRUFileCache::DeletePhyFile(const std::string& path) {
  std::string fullpath = path::Join(workdir_, path);
  FileInfo fi;
  Status s = StatPath(fullpath, &fi);
  if (!s.ok()) {
    LOG_ERROR << "failed to stat file: " << fullpath << ": " << s;
    return false;
  }

  if (fi.not_exists()) return true;

  s = DeleteFile(fullpath);
  if (!s.ok()) {
    LOG_ERROR << "failed to delete file: " << fullpath << ": " << s;
    return false;
  }

  return true;
}

void LRUFileCache::CreateLogTable() {
  db_->RunSql(
      "create table if not exists file_cache_log( "
      " id varchar primary key, "
      " path varchar not null, "
      " file_size integer, "
      " src_digest varchar, "
      " aes_digest varchar, "
      " state integer default 0, "
      " sticky integer default 0, "
      " lru_time integer default 0)");

  db_->RunSql(
      "create index if not exists file_cache_log_idx_time "
      " on file_cache_log(lru_time)");
}

bool LRUFileCache::LoadLogEntry(db::StmtPtr& stmt, FileLogEntry* e) {
  assert(!stmt->eof());
  int p = 0;
  e->id = stmt->GetText(p++);
  e->path = stmt->GetText(p++);
  e->file_size = stmt->GetInt64(p++);
  e->src_digest = stmt->GetText(p++);
  e->aes_digest = stmt->GetText(p++);
  e->state = stmt->GetInt32(p++);
  e->sticky = stmt->GetInt32(p++);
  e->lru_time = Timestamp(stmt->GetInt64(p++));
  return true;
}

bool LRUFileCache::UpdateLogEntry(const FileLogEntry& e) {
  db::StmtPtr stmt = db_->Prepare(
      "replace into file_cache_log "
      "(id, path, file_size, src_digest, aes_digest, state, sticky, lru_time) "
      "values "
      "(?, ?, ?, ?, ?, ?, ?, ?)");
  stmt->BindText(e.id);
  stmt->BindText(e.path);
  stmt->BindInt64(e.file_size);
  stmt->BindText(e.src_digest);
  stmt->BindText(e.aes_digest);
  stmt->BindInt32(e.state);
  stmt->BindInt32(e.sticky);
  stmt->BindInt64(e.lru_time.MicroSecondsSinceEpoch());
  stmt->Step();
  return true;
}

bool LRUFileCache::LogStartMoving(const std::string& id,
                                  const std::string& path, size_t file_size) {
  FileLogEntry e;
  e.id = id;
  e.path = path;
  e.file_size = file_size;
  e.state = kFileMoving;
  e.sticky = Sticky::removable;
  e.lru_time = Timestamp::Now();
  return UpdateLogEntry(e);
}

bool LRUFileCache::LogMovingFailed(const std::string& id,
                                   const std::string& path, size_t file_size) {
  FileLogEntry e;
  e.id = id;
  e.path = path;
  e.file_size = file_size;
  e.state = kFileFailed;
  e.sticky = Sticky::removable;
  e.lru_time = Timestamp::Now();
  return UpdateLogEntry(e);
}

bool LRUFileCache::LogMovingResult(const std::string& id,
                                   const std::string& path, size_t file_size,
                                   const std::string& src_digest,
                                   const std::string& aes_digest, Sticky sticky,
                                   bool success) {
  // assert (!IdExists(id));
  FileLogEntry e;
  e.id = id;
  e.path = path;
  e.file_size = file_size;
  e.src_digest = src_digest;
  e.aes_digest = aes_digest;
  e.state = (success ? kFileNormal : kFileFailed);
  e.sticky = sticky;
  e.lru_time = Timestamp::Now();
  return UpdateLogEntry(e);
}

bool LRUFileCache::DeleteFileCacheLog(const std::string& id) {
  db::StmtPtr stmt = db_->Prepare("delete from file_cache_log where id = ?");
  stmt->BindText(id);
  stmt->Step();
  return true;
}

bool LRUFileCache::IdExists(const std::string& id) {
  db::StmtPtr stmt = db_->Prepare("select id from file_cache_log where id = ?");
  stmt->BindText(id);
  stmt->Step();
  return !stmt->eof();
}

void LRUFileCache::RemoveFilesWithIncompleteLog() {
  db::StmtPtr stmt =
      db_->Prepare("select * from file_cache_log where state != ?");
  stmt->BindInt32(kFileNormal);
  stmt->Step();

  std::vector<FileLogEntry> errors;
  for (; !stmt->eof(); stmt->Step()) {
    FileLogEntry e;
    LoadLogEntry(stmt, &e);
    errors.push_back(e);
  }

  stmt.reset();
  for (size_t i = 0; i < errors.size(); ++i) {
    FileLogEntry& e = errors[i];
    if (DeletePhyFile(e.path) == false) {
      LOG_ERROR << "failed to delete incomplete cache file: " << e.path;
      continue;
    }
    DeleteFileCacheLog(e.id);
  }
}

void LRUFileCache::LoadCacheFiles() {
  assert(used_ == 0);
  db::StmtPtr stmt = db_->Prepare(
      "select * from file_cache_log where state = ?"
      " order by lru_time ASC");
  stmt->BindInt32(kFileNormal);
  stmt->Step();

  for (; !stmt->eof(); stmt->Step()) {
    FileLogEntry e;
    LoadLogEntry(stmt, &e);
    LRUFileHandle* h = new LRUFileHandle();
    h->id = e.id;
    h->path = e.path;
    h->size = e.file_size;
    h->sticky = static_cast<Sticky>(e.sticky);
    h->src_digest = e.src_digest;
    h->aes_digest = e.aes_digest;
    h->refs = 1;

    h->InsertBefore(files_.begin());
    used_ += h->size;

    LOG_INFO << "InsertHandle: " << h->id << " : " << h->path << " : "
             << h->size << "-" << used_;
    table_.insert(std::make_pair(h->id, h));
  }
}

bool LRUFileCache::MakeSpace(size_t size) {
  LinkNode<LRUFileHandle>* n = files_.rbegin();

  while (used_ + static_cast<int64_t>(size) > capacity_ && n != files_.rend()) {
    LinkNode<LRUFileHandle>* cur = n;
    n = cur->prev();

    LRUFileHandle* h = cur->value();
    if (h->sticky == Sticky::sticky) {
      continue;
    } else {
      cur->RemoveFromList();
      table_.erase(h->id);
      LOG_INFO << "Remove file from cache: " << h->id << " : " << h->path;
      UnrefHandle(h);
    }
  }

  return (used_ + static_cast<int64_t>(size) <= capacity_);
}

void LRUFileCache::UnrefHandle(LRUFileHandle* h) {
  assert(h->refs > 0);
  --h->refs;

  LOG_DEBUG << "UnrefHandle: " << h->id;
  if (h->refs == 0) {
    LOG_WARN << "DestroyHandle: " << h->id << " : " << h->path << " : "
             << h->size << "-" << used_;
    DeleteFileCacheLog(h->id);
    used_ -= h->size;
    assert(used_ >= 0);
    h->RemoveFromList();
    if (DeletePhyFile(h->path) == false) {
      LogMovingFailed(h->id, h->path, h->size);
      LOG_ERROR << "failed to delete incomplete cache file: " << h->path;
    }
    delete h;
  }
}

void LRUFileCache::MarkAsRemovable(const std::string& id) {
  LockGuard g(mu_);
  IdMap::iterator it = table_.find(id);
  if (it == table_.end()) {
    LOG_ERROR << "try to mark file removable but not found: " << id;
    return;
  }

  LRUFileHandle* h = it->second;
  assert(h != NULL && h->id == id);
  h->sticky = Sticky::removable;
  LOG_INFO << "Mark file as removable: " << h->id << " : " << h->path;

  FileLogEntry e;
  e.id = h->id;
  e.path = h->path;
  e.file_size = h->size;
  e.src_digest = h->src_digest;
  e.aes_digest = h->aes_digest;
  e.sticky = Sticky::removable;
  e.state = kFileNormal;
  e.lru_time = Timestamp::Now();
  UpdateLogEntry(e);

  ::chmod(path::Join(workdir_, h->path).c_str(), 0600);
}

void LRUFileCache::Release(Handle* handle) {
  if (handle == NULL) return;
  LockGuard g(mu_);
  LRUFileHandle* h = reinterpret_cast<LRUFileHandle*>(handle);
  UnrefHandle(h);
}

std::string LRUFileCache::FilePath(Handle* handle) {
  LRUFileHandle* h = reinterpret_cast<LRUFileHandle*>(handle);
  return path::Join(workdir_, h->path);
}

bool LRUFileCache::GetFileDetail(Handle* handle, FileDetail* detail) {
  LRUFileHandle* h = reinterpret_cast<LRUFileHandle*>(handle);

  detail->id = h->id;
  detail->path = path::Join(workdir_, h->path);
  detail->src_digest = h->src_digest;
  detail->aes_digest = h->aes_digest;
  return true;
}

void LRUFileCache::Erase(const std::string& id) {
  LOG_DEBUG << "EraseHandle: " << id;

  LockGuard g(mu_);
  IdMap::iterator it = table_.find(id);
  if (it == table_.end()) return;

  LRUFileHandle* h = it->second;
  assert(h != NULL);
  table_.erase(id);
  UnrefHandle(h);
}

LRUFileCache::Handle* LRUFileCache::InsertFromTmpFile(
    const std::string& id, const std::string& path, size_t alloc_tmp_size,
    const std::string& src_digest, const std::string& aes_digest,
    Sticky sticky) {
  LOG_DEBUG << "InsertHandle: " << id;

  LRUFileHandle* h = NULL;
  Status s;
  LockGuard g(mu_);

  FileInfo fi;
  s = StatPath(path, &fi);
  if (!s.ok() || !fi.is_regular()) return NULL;

  if (fi.size > alloc_tmp_size &&
      MakeSpace(fi.size - alloc_tmp_size) == false) {
    LOG_ERROR << "local cache is full of sticky files: " << this->capacity_
              << " - " << this->used_;
    return NULL;
  }

  // delete old one
  if (table_.find(id) != table_.end()) {
    h = table_[id];
    table_.erase(id);
    UnrefHandle(h);
    h = NULL;
  }

  // write log before moving
  std::string dest = GenFileId();
  LogStartMoving(id, dest, fi.size);

  // try to move file
  s = dicom::RenameFile(path, path::Join(workdir_, dest));
  if (!s.ok()) {
    LogMovingFailed(id, dest, fi.size);
    return NULL;
  }

  LogMovingResult(id, dest, fi.size, src_digest, aes_digest, Sticky::sticky,
                  true);
  h = new LRUFileHandle();
  h->id = id;
  h->path = dest;
  h->size = fi.size;
  h->sticky = sticky;
  h->src_digest = src_digest;
  h->aes_digest = aes_digest;
  h->refs = 1;

  tmp_files_.erase(path);
  used_ -= alloc_tmp_size;
  used_ += fi.size;

  h->InsertBefore(files_.begin());
  LOG_INFO << "InsertHandle: " << h->id << " : " << h->path << " : " << h->size
           << "-" << used_;

  table_.insert(std::make_pair(id, h));

  if (sticky == Sticky::removable)
    ::chmod(path::Join(workdir_, h->path).c_str(), 0600);

  ++h->refs;
  return reinterpret_cast<LRUFileCache::Handle*>(h);
}

LRUFileCache::Handle* LRUFileCache::Insert(
    const std::string& id, const std::string& path, size_t file_size,
    const std::string& src_digest, const std::string& aes_digest,
    FileInTmpDir already_in_tmp_dir, Sticky sticky) {
  LOG_DEBUG << "InsertHandle: " << id;

  LRUFileHandle* h = NULL;
  LockGuard g(mu_);

  if (already_in_tmp_dir == FileInTmpDir::no && MakeSpace(file_size) == false) {
    return NULL;
  }

  // delete old one
  if (table_.find(id) != table_.end()) {
    h = table_[id];
    table_.erase(id);
    UnrefHandle(h);
    h = NULL;
  }

  // write log before moving
  std::string dest = GenFileId();
  LogStartMoving(id, dest, file_size);

  // try to move file
  Status s = dicom::RenameFile(path, path::Join(workdir_, dest));
  if (!s.ok()) {
    LogMovingFailed(id, dest, file_size);
    return NULL;
  }

  LogMovingResult(id, dest, file_size, src_digest, aes_digest, sticky, true);
  h = new LRUFileHandle();
  h->id = id;
  h->path = dest;
  h->size = file_size;
  h->sticky = sticky;
  h->src_digest = src_digest;
  h->aes_digest = aes_digest;
  h->refs = 1;

  if (!already_in_tmp_dir) {
    used_ += h->size;
  } else {
    tmp_files_.erase(path);
  }

  h->InsertBefore(files_.begin());
  LOG_INFO << "InsertHandle: " << h->id << " : " << h->path << " : " << h->size
           << "-" << used_;

  table_.insert(std::make_pair(id, h));

  ++h->refs;
  return reinterpret_cast<LRUFileCache::Handle*>(h);
}

LRUFileCache::Handle* LRUFileCache::Lookup(const std::string& id) {
  LockGuard g(mu_);
  IdMap::iterator it = table_.find(id);
  if (it == table_.end()) return NULL;

  LRUFileHandle* h = it->second;
  assert(h != NULL);

  h->RemoveFromList();
  h->InsertBefore(files_.begin());

  FileLogEntry e;
  e.id = h->id;
  e.path = h->path;
  e.file_size = h->size;
  e.src_digest = h->src_digest;
  e.aes_digest = h->aes_digest;
  e.sticky = h->sticky;
  e.state = kFileNormal;
  e.lru_time = Timestamp::Now();
  UpdateLogEntry(e);

  h->refs++;
  // assert file exists
  FileInfo fi;
  StatPath(FilePath(reinterpret_cast<Handle*>(h)), &fi);
  if (!fi.is_regular()) {
    LOG_ERROR << "CACHE ERROR: file missing: " << h->id;
  }

  LOG_DEBUG << "LookupHandle: " << h->id;
  return reinterpret_cast<Handle*>(h);
}

bool LRUFileCache::InitCache() {
  dicom::CreateDir(workdir_);

  Status s;
  FileInfo fi;
  s = dicom::StatPath(workdir_, &fi);
  if (!s.ok() || !fi.is_dir()) {
    LOG_ERROR << "failed to create working directory";
    return false;
  }

  tmp_dir_ = path::Join(workdir_, "tmp");
  s = dicom::CreateDir(tmp_dir_);
  if (!s.ok()) {
    LOG_ERROR << "failed to create working directory";
    return false;
  }

  s = db::Database::Open(path::Join(workdir_, "manifest.db"), &db_);
  if (!s.ok()) {
    LOG_ERROR << "failed to create cache database";
    return false;
  }

  this->CreateLogTable();
  this->DeleteAllTempFiles();
  this->RemoveFilesWithIncompleteLog();
  this->LoadCacheFiles();
  // Remove file if user change the configuration to
  // use small local storage space
  this->MakeSpace(1);
  return true;
}

LRUFileCache* NewLocalCache(const std::string& workdir, size_t capacity) {
  LRUFileCache* c = new LRUFileCache(workdir, capacity);
  if (!c->InitCache()) {
    delete c;
    return NULL;
  }

  return c;
}

void LRUFileCache::DeleteAllTempFiles() {
  DIR* d = ::opendir(tmp_dir_.c_str());
  if (d == NULL) return;

  struct dirent* dp = NULL;

  do {
    if ((dp = ::readdir(d)) != NULL) {
      if (strcmp(dp->d_name, ".") == 0 || strcmp(dp->d_name, "..") == 0) {
        continue;
      }

      Status s = DeleteFile(path::Join(tmp_dir_, dp->d_name));
      if (!s.ok()) LOG_ERROR << "failed to delete tmp file: " << dp->d_name;
    }
  } while (dp != NULL);

  ::closedir(d);
}

bool LRUFileCache::AllocTmpFile(const char* domain, size_t size,
                                std::string* path) {
  LockGuard g(mu_);

  if (size > 0) {
    MakeSpace(size);
    if (used_ + static_cast<int64_t>(size) > capacity_) {
      LOG_ERROR << "No space for tmp file of size: " << size;
      return false;
    }
  }

  std::string uid = GenFileId();
  (*path) = StringPrintf("%s/%s.tmp", tmp_dir_.c_str(), uid.c_str());
  used_ += size;
  tmp_files_.insert(std::make_pair((*path), size));
  LOG_INFO << domain << ": AllocTmpFile: " << (*path) << " : " << size;
  return true;
}

void LRUFileCache::ReleaseTmpFile(const char* domain, const std::string& path) {
  LockGuard g(mu_);
  if (tmp_files_.find(path) == tmp_files_.end()) {
    LOG_ERROR << "try to release non-existing tmp file: " << path;
    return;
  }

  size_t size = tmp_files_.find(path)->second;
  LOG_INFO << domain << ": ReleaseTmpFile: " << path << " : " << size;
  assert(static_cast<int64_t>(size) <= used_);
  used_ -= size;
  assert(used_ >= 0);

  FileInfo fi;
  Status s = StatPath(path, &fi);
  if (fi.not_exists()) {
    return;
  }

  s = DeleteFile(path);
  if (!s.ok()) {
    LOG_ERROR << "failed to delete cache file: " << path << ": " << s;
  }
}

void LRUFileCache::GetStat(CacheStat* stat) {
  LockGuard g(mu_);
  stat->total_space = this->capacity_;
  stat->used_space = this->used_;
  stat->removable_space = 0;

  LinkNode<LRUFileHandle>* n = files_.rbegin();
  while (n != files_.rend()) {
    LRUFileHandle* h = n->value();
    if (!h->sticky) stat->removable_space += h->size;
    n = n->prev();
  }
}

LRUFileCache::Handle* SaveFileToCacheWithAES(LRUFileCache* cache,
                                             const std::string& path,
                                             const std::string& fid,
                                             std::string hex_key,
                                             HashEntry* hash) {
  Status s;
  FileInfo src_fi;
  FileReader in;
  FileWriter out;
  Encryptor encryptor;
  std::string tmpfile;
  const char* domain = "AES_SAVEER";

  s = StatPath(path, &src_fi);
  if (!s.ok() || !src_fi.is_regular()) {
    LOG_ERROR << "source file is not regular file: " << path;
    return NULL;
  }

  s = in.Open(path);
  if (!s.ok()) {
    LOG_ERROR << "failed to open source file: " << path;
    return NULL;
  }

  size_t alloc_tmp_size = (src_fi.size + 16) / 16 * 16;
  if (cache->AllocTmpFile(domain, alloc_tmp_size, &tmpfile) == false) {
    LOG_ERROR << "failed to alloc temp file of size: " << alloc_tmp_size;
    return NULL;
  }

  if (!out.Open(tmpfile).ok()) {
    cache->ReleaseTmpFile(domain, tmpfile);
    LOG_ERROR << "failed to open temp file: " << tmpfile;
    return NULL;
  }

  s = encryptor.DoEncryt(hex_key, &in, &out, hash);
  in.Close();
  out.Close();

  if (!s.ok()) {
    cache->ReleaseTmpFile(domain, tmpfile);
    LOG_ERROR << "failed to encrypt file: " << fid;
    return NULL;
  }

  LRUFileCache::Handle* h =
      cache->InsertFromTmpFile(fid, tmpfile, alloc_tmp_size, hash->src_md5,
                               hash->enc_digest, Sticky::sticky);
  if (h == NULL) {
    // cache->Insert() will remove temp file in success
    LOG_ERROR << "failed to insert file to cache: " << path;
    cache->ReleaseTmpFile(domain, tmpfile);
  }

  return h;
}
}
